#!/usr/bin/env python3
"""
ComfyUI 环境配置脚本
统一管理所有路径、sys.path、PYTHONPATH等环境变量
提供全局环境路径供其他脚本使用
"""

import sys
import os
import platform
from pathlib import Path
import subprocess

# ==================== 全局路径配置 ====================

# 获取py_scripts目录（当前文件在 py_scripts/public/ 下）
BASE_DIR = Path(__file__).resolve().parent.parent  # 指向 py_scripts

# 项目根目录（从py_scripts往上2级）
PROJECT_ROOT = BASE_DIR.parent.parent

# 定义重要目录
SCRIPTS_PUBLIC = BASE_DIR / 'public'  # 当前目录
PY_SCRIPTS_DIR = BASE_DIR  # 工作目录
COMFYUI_DIR = PROJECT_ROOT / 'ComfyUI'  # ComfyUI主程序目录
COMFYUI_CUSTOM_NODES = COMFYUI_DIR / 'custom_nodes'  # ComfyUI自定义节点
CONFIG_CUSTOM_NODES = PROJECT_ROOT / 'Config' / 'custom_nodes'  # 配置自定义节点
PYTHON_DIR = PROJECT_ROOT / 'python'  # Python环境目录
PYTHON_SCRIPTS = PYTHON_DIR / 'Scripts'  # Windows Python脚本目录
PYTHON_LIB = PYTHON_DIR / 'Lib' / 'site-packages'  # Python包目录
CONFIG_DIR = BASE_DIR / 'config'  # 配置目录

# 统一的下载相关全局路径
DOWNMODELS_DIR = PROJECT_ROOT / 'downModels'   # HuggingFace/批量下载保存目录
URLDOWN_DIR = PROJECT_ROOT / 'URLDown'         # URL下载保存目录
BATCHDOWN_DIR = PROJECT_ROOT / 'downmodels'    # 批量下载保存目录（兼容小写）

# 根据操作系统设置Python可执行文件路径
if platform.system() == "Windows":
    PYTHON_EXE = PYTHON_DIR / "python.exe"
    PYTHON_CMD = str(PYTHON_EXE)
else:
    PYTHON_EXE = PYTHON_DIR / "bin" / "python3"
    PYTHON_CMD = str(PYTHON_EXE)

# ==================== 环境设置函数 ====================

def ensure_package(pkg_name, import_name=None):
    """
    检查并自动安装指定包
    :param pkg_name: pip 包名
    :param import_name: import 时的模块名（如不同于pip包名）
    """
    import_name = import_name or pkg_name
    try:
        __import__(import_name)
    except ImportError:
        print(f"未检测到依赖包 {pkg_name}，正在自动安装...")
        subprocess.check_call([sys.executable, "-m", "pip", "install", pkg_name])
        print(f"已安装 {pkg_name}")

def setup_environment():
    """设置统一的环境配置"""
    
    # 需要添加到sys.path的目录列表
    paths_to_add = [
        str(PY_SCRIPTS_DIR),      # 工作目录
    ]
    
    # 添加到sys.path（避免重复添加）
    for path in paths_to_add:
        if path not in sys.path:
            sys.path.insert(0, path)
    
    # 设置PYTHONPATH环境变量
    pythonpath_parts = [str(SCRIPTS_PUBLIC), str(COMFYUI_DIR), str(COMFYUI_CUSTOM_NODES)]
    os.environ['PYTHONPATH'] = os.pathsep.join(pythonpath_parts)
    
    # 设置其他有用的环境变量
    os.environ['COMFYUI_BASE_DIR'] = str(BASE_DIR)
    os.environ['COMFYUI_SCRIPTS_DIR'] = str(SCRIPTS_PUBLIC)
    os.environ['COMFYUI_CUSTOM_NODES_DIR'] = str(COMFYUI_CUSTOM_NODES)
    os.environ['COMFYUI_PROJECT_ROOT'] = str(PROJECT_ROOT)
    
    # 返回配置信息
    config = {
        'BASE_DIR': BASE_DIR,
        'PROJECT_ROOT': PROJECT_ROOT,
        'SCRIPTS_PUBLIC': SCRIPTS_PUBLIC,
        'COMFYUI_DIR': COMFYUI_DIR,
        'COMFYUI_CUSTOM_NODES': COMFYUI_CUSTOM_NODES,
        'CONFIG_CUSTOM_NODES': CONFIG_CUSTOM_NODES,
        'PYTHON_DIR': PYTHON_DIR,
        'PYTHON_EXE': PYTHON_EXE,
        'PYTHON_CMD': PYTHON_CMD,
        'PYTHON_SCRIPTS': PYTHON_SCRIPTS,
        'PYTHON_LIB': PYTHON_LIB,
        'CONFIG_DIR': CONFIG_DIR,
        'paths_added': paths_to_add,
        'DOWNMODELS_DIR': DOWNMODELS_DIR,
        'URLDOWN_DIR': URLDOWN_DIR,
        'BATCHDOWN_DIR': BATCHDOWN_DIR
    }
    
    return config

# ==================== 全局环境变量 ====================

# 自动设置环境
ENV_CONFIG = setup_environment()

# 为了方便使用，直接导出主要路径
# 其他脚本可以直接使用这些变量
PROJECT_ROOT = ENV_CONFIG['PROJECT_ROOT']
COMFYUI_DIR = ENV_CONFIG['COMFYUI_DIR']
PYTHON_DIR = ENV_CONFIG['PYTHON_DIR']
PYTHON_EXE = ENV_CONFIG['PYTHON_EXE']
PYTHON_CMD = ENV_CONFIG['PYTHON_CMD']
CONFIG_DIR = ENV_CONFIG['CONFIG_DIR']

# ==================== 工具函数 ====================

def print_environment_info(config=None):
    """打印环境配置信息"""
    if config is None:
        config = ENV_CONFIG
    
    print("=" * 60)
    print("           ComfyUI 环境配置信息")
    print("=" * 60)
    print(f"项目根目录: {config['PROJECT_ROOT']}")
    print(f"py_scripts目录: {config['BASE_DIR']}")
    print(f"脚本目录: {config['SCRIPTS_PUBLIC']}")
    print(f"ComfyUI目录: {config['COMFYUI_DIR']}")
    print(f"自定义节点目录: {config['COMFYUI_CUSTOM_NODES']}")
    print(f"配置节点目录: {config['CONFIG_CUSTOM_NODES']}")
    print(f"Python目录: {config['PYTHON_DIR']}")
    print(f"Python可执行文件: {config['PYTHON_EXE']}")
    print(f"Python包目录: {config['PYTHON_LIB']}")
    print(f"PYTHONPATH: {os.environ.get('PYTHONPATH', '未设置')}")
    print(f"已添加到sys.path的目录数量: {len(config['paths_added'])}")
    print(f"HuggingFace下载目录: {config['DOWNMODELS_DIR']}")
    print(f"URL下载目录: {config['URLDOWN_DIR']}")
    print(f"批量下载目录: {config['BATCHDOWN_DIR']}")
    print("=" * 60)

def check_paths_exist(config=None):
    """检查重要路径是否存在"""
    if config is None:
        config = ENV_CONFIG
    
    missing_paths = []
    for name, path in config.items():
        if name.endswith('_DIR') or name.endswith('_NODES') or name.endswith('_PUBLIC'):
            if not path.exists():
                missing_paths.append((name, path))
    
    if missing_paths:
        print("⚠️  警告：以下目录不存在：")
        for name, path in missing_paths:
            print(f"  {name}: {path}")
        return False
    else:
        print("✅ 所有重要目录都存在")
        return True

def get_python_cmd():
    """获取Python命令字符串"""
    return PYTHON_CMD

def get_project_root():
    """获取项目根目录"""
    return PROJECT_ROOT

def get_comfyui_dir():
    """获取ComfyUI目录"""
    return COMFYUI_DIR

def get_python_dir():
    """获取Python目录"""
    return PYTHON_DIR

# ==================== 依赖包管理 ====================

# 依赖包列表
REQUIRED_PACKAGES = [
    ("requests", None),
    ("huggingface_hub", None),
    ("packaging", None),
    ("setuptools", None),
    ("psutil", None),
    # PyQt6只在Windows下需要，Linux下可选
    # ("PyQt6", None) if platform.system() == "Windows" else None,
]
# 过滤掉None值
REQUIRED_PACKAGES = [pkg for pkg in REQUIRED_PACKAGES if pkg is not None]

def install_required_packages():
    """安装必需的依赖包"""
    print("📦 检查并安装必需的依赖包...")
    for pkg, import_name in REQUIRED_PACKAGES:
        ensure_package(pkg, import_name)
    print("✅ 依赖包检查完成")

# ==================== 主函数 ====================

def main():
    """命令行入口点"""
    print("🔧 ComfyUI 环境配置工具")
    print("=" * 60)
    
    # 安装依赖包
    install_required_packages()
    
    # 打印环境信息
    print_environment_info()
    
    # 检查路径
    check_paths_exist()
    
    print("\n💡 使用提示:")
    print("  其他脚本可以直接导入以下变量:")
    print("  from env_setup import PROJECT_ROOT, COMFYUI_DIR, PYTHON_CMD")
    print("  from env_setup import get_python_cmd, get_project_root")

if __name__ == "__main__":
    main() 